# -*- coding: utf-8 -*-
"""
Created on Fri Sep 28 17:18:19 2018

@author: HatLabUnderGrad2
"""

# fitting library for the hist plot 

import numpy as np
import matplotlib.pyplot as plt
import scipy.optimize as opt

def twoD_Gaussian(xy, amplitude, xo, yo, sigma_x, sigma_y):
    '''
    2D Gaussian function for fitting. 
    
    Input:
        (x,y): the x and y variable for the 2D gaussian. x and y are generated from np.meshgrid(x, y).
        amplitude (double): the amplitude of the gaussian (usually the max count of the hist data). 
        xo, yo (double): the center of the gaussian. 
        sigma_x, sigma_y (double): the sigma of the gaussian in x and y direction. 
        theta (double): the angle between the axis of the gaussian and x axis
        offset (double): the offset of the gaussian compare to 0 at infinity
        
    Output:
        the data for the 2D gaussian in a 1D array
    '''
#    print ('use this guassian')
    theta = 0
    offset = 0
    xo = float(xo)
    yo = float(yo)    
    (x, y) = xy
    a = (np.cos(theta)**2)/(2*sigma_x**2) + (np.sin(theta)**2)/(2*sigma_y**2)
    b = -(np.sin(2*theta))/(4*sigma_x**2) + (np.sin(2*theta))/(4*sigma_y**2)
    c = (np.sin(theta)**2)/(2*sigma_x**2) + (np.cos(theta)**2)/(2*sigma_y**2)
    g = offset + amplitude*np.exp( - (a*((x-xo)**2) + 2*b*(x-xo)*(y-yo)
                            + c*((y-yo)**2)))
    return g.ravel()



def oneD_Gaussian(x, amplitude, xo, sigma_x):
    '''
    2D Gaussian function for fitting. 
    
    Input:
        (x,y): the x and y variable for the 2D gaussian. x and y are generated from np.meshgrid(x, y).
        amplitude (double): the amplitude of the gaussian (usually the max count of the hist data). 
        xo, yo (double): the center of the gaussian. 
        sigma_x, sigma_y (double): the sigma of the gaussian in x and y direction. 
        theta (double): the angle between the axis of the gaussian and x axis
        offset (double): the offset of the gaussian compare to 0 at infinity
        
    Output:
        the data for the 2D gaussian in a 1D array
    '''
#    print ('use this guassian')
#    theta = 0
    offset = 0
    xo = float(xo)
#    yo = float(yo)    
#    a = (np.cos(theta)**2)/(2*sigma_x**2) + (np.sin(theta)**2)/(2*sigma_y**2)
#    b = -(np.sin(2*theta))/(4*sigma_x**2) + (np.sin(2*theta))/(4*sigma_y**2)
#    c = (np.sin(theta)**2)/(2*sigma_x**2) + (np.cos(theta)**2)/(2*sigma_y**2)
#    g = offset + amplitude*np.exp( - (a*((x-xo)**2) + 2*b*(x-xo)*(y-yo)
#                            + c*((y-yo)**2)))
    xo = float(xo)
    g = offset + amplitude*np.exp(-(x - xo)**2/(2*sigma_x**2))
    return g.ravel()


  
def two_blob(xy, amp1, amp2, xo1, yo1, xo2, yo2, sigma_x_1, sigma_y_1, sigma_x_2, sigma_y_2):
    '''
    Function for fitting two 2D Gaussian in the same data. 
    
    Input:
        (x,y): the x and y variable for the 2D gaussian. x and y are generated from np.meshgrid(x, y).
        amp1, amp2 (double): the amplitude of each gaussian (usually the max count of the hist data). 
        xo1, yo1, xo2, yo2 (double): the center of each gaussian. 
        sigma_x_1, sigma_y_1, sigma_x_2, sigma_y_2 (double): the sigma of each gaussian in x and y direction. 
        theta (double): the angle between the axis of the gaussian and x axis
        offset (double): the offset of the gaussian compare to 0 at infinity
        
    Output:
        the data for the two 2D gaussian in a 1D array    
    
    '''
    return twoD_Gaussian(xy, amp1, xo1, yo1, sigma_x_1, sigma_y_1) \
            + twoD_Gaussian(xy, amp2, xo2, yo2, sigma_x_2, sigma_y_2)


def three_blob(xy, amp1, amp2, amp3, xo1, yo1, xo2, yo2, xo3, yo3, sigma_x_1, sigma_y_1, sigma_x_2, sigma_y_2, sigma_x_3, sigma_y_3):
    '''
    Function for fitting two 2D Gaussian in the same data. 
    
    Input:
        (x,y): the x and y variable for the 2D gaussian. x and y are generated from np.meshgrid(x, y).
        amp1, amp2 (double): the amplitude of each gaussian (usually the max count of the hist data). 
        xo1, yo1, xo2, yo2 (double): the center of each gaussian. 
        sigma_x_1, sigma_y_1, sigma_x_2, sigma_y_2 (double): the sigma of each gaussian in x and y direction. 
        theta (double): the angle between the axis of the gaussian and x axis
        offset (double): the offset of the gaussian compare to 0 at infinity
        
    Output:
        the data for the two 2D gaussian in a 1D array    
    
    '''
    return twoD_Gaussian(xy, amp1, xo1, yo1, sigma_x_1, sigma_y_1) \
            + twoD_Gaussian(xy, amp2, xo2, yo2, sigma_x_2, sigma_y_2) \
            + twoD_Gaussian(xy, amp3, xo3, yo3, sigma_x_3, sigma_y_3)




def fit_2D_Gaussian(xy, data, initial_guess, fit_model, plot=False, bounds=None):
    '''
    Fit the data to a 2D gaussian model.
    
    Input:
        (x, y): the x and y variable for the 2D gaussian. x and y are generated from np.meshgrid(x, y).
        initial_guess (tuple): initial guess for the fitting, depends on the fit model, this could be different
        fit_model (function): the model that this data is fitting to. 
        plot (boolean): set to True if want to plot the fitting result. Defalut to False.
    
    Ouput:
        fit_result (np array): the paraemters from the best fitting result. The order of the parameters are the same as the initial guess.
    '''
#    print ('using it')
    if bounds is not None:
        print ('Bounds!')
        fit_bounds = bounds
    else:
        if fit_model == two_blob:
            fit_bounds = ((0, 0, -np.inf, -np.inf, -np.inf, -np.inf, 0, 0, 0, 0), np.inf)
        elif fit_model == twoD_Gaussian:
            fit_bounds = ((0, -np.inf, -np.inf, 0, 0), np.inf)
        elif fit_model == three_blob:
            fit_bounds = ((0, 0, 0, -np.inf, -np.inf, -np.inf, -np.inf, -np.inf, -np.inf, 0, 0, 0, 0, 0, 0), np.inf)
        else:
            print ('No such fit model.')
        

#    try:
    fit_result, pcov = opt.curve_fit(fit_model, xy, data, bounds = fit_bounds, p0=initial_guess, maxfev=5000)
#    except:
#        raise NameError('Fitting wrong')
        
    if plot:
        plot_fit_result(xy, data, fit_result, fit_model)
    return fit_result


def fit_1D_Gaussian(x, data, initial_guess, fit_model, plot=False, bounds=None):
    '''
    Fit the data to a 2D gaussian model.
    
    Input:
        (x, y): the x and y variable for the 2D gaussian. x and y are generated from np.meshgrid(x, y).
        initial_guess (tuple): initial guess for the fitting, depends on the fit model, this could be different
        fit_model (function): the model that this data is fitting to. 
        plot (boolean): set to True if want to plot the fitting result. Defalut to False.
    
    Ouput:
        fit_result (np array): the paraemters from the best fitting result. The order of the parameters are the same as the initial guess.
    '''
#    print ('using it')
    if bounds is not None:
        print ('Bounds!')
        fit_bounds = bounds
    else:
        if fit_model == oneD_Gaussian:
            fit_bounds = ((0, -np.inf, 0), np.inf)
        else:
            print ('No such fit model.')
        

#    try:
    fit_result, pcov = opt.curve_fit(fit_model, x, data, bounds = fit_bounds, p0=initial_guess, maxfev=5000)
#    except:
#        raise NameError('Fitting wrong')
        
    if plot:
        plot_1D_fit_result(x, data, fit_result, fit_model)
    return fit_result

def plot_fit_result(xy, data, fit_result, fit_model):
    '''
    Plot the data and its fitting result.

    Input: 
        (x, y): the x and y variable for the 2D gaussian. x and y are generated from np.meshgrid(x, y).
        data (np array): the data that is used for the fitting. 
        fit_result (np array): the paraemters from the best fitting result. The order of the parameters are the same as the initial guess.
        fit_model (function): the model that this data is fitting to.
    
    Output:
        None
    
    '''
    fig, ax = plt.subplots(1, 1)    
    (x, y) = xy
    data_fitted = fit_model(xy, *fit_result)
#    ax.hold(True)
    ax.imshow(data.reshape(y.shape[0], x.shape[1]), cmap=plt.cm.jet, origin='lower',
        extent=(x.min(), x.max(), y.min(), y.max()))
    ax.contour(x, y, data_fitted.reshape(y.shape[0], x.shape[1]), 8, colors='w')
    plt.show()
    

def plot_1D_fit_result(x, data, fit_result, fit_model):
    '''
    Plot the data and its fitting result.

    Input: 
        (x, y): the x and y variable for the 2D gaussian. x and y are generated from np.meshgrid(x, y).
        data (np array): the data that is used for the fitting. 
        fit_result (np array): the paraemters from the best fitting result. The order of the parameters are the same as the initial guess.
        fit_model (function): the model that this data is fitting to.
    
    Output:
        None
    
    '''
    fig, ax = plt.subplots(1, 1)    
    data_fitted = fit_model(x, *fit_result)
#    ax.hold(True)
#    ax.imshow(data.reshape(y.shape[0], x.shape[1]), cmap=plt.cm.jet, origin='bottom',
#        extent=(x.min(), x.max(), y.min(), y.max()))
#    ax.contour(x, y, data_fitted.reshape(y.shape[0], x.shape[1]), 8, colors='w')
    ax.plot(x, data, '*')
    ax.plot(x, data_fitted, '-')
    plt.show()


if __name__ == '__main__':

    ### create x and y axis
    x = np.linspace(-200, 200, 201)
    y = np.linspace(-200, 200, 201)
#    x, y = np.meshgrid(x, y)
    xy = np.meshgrid(x, y)
    
    ### create data
#    data = twoD_Gaussian((x, y), 3, 100, 100, 20, 40)#, 0*np.pi/180.0, 10)
#    data2 = twoD_Gaussian((x, y), 3, -100, -100, 20, 40)#, 0*np.pi/180.0, 10)
    
    data = twoD_Gaussian(xy, 3, 100, 100, 20, 40)
    data2 = twoD_Gaussian(xy, 3, -100, -100, 20, 40)
    data3 = twoD_Gaussian(xy, 3, -100, 100, 20, 40)
    
    ### add noise to data
    initial_guess = (3, 100, 100, 20, 40)#, 0, 10)
    data_noisy = data + 0.2*np.random.normal(size=data.shape)
    
    initial_guess2 = (3, 3, 100, 100, -100, -100, 20, 40, 20, 40)#, 0, 10)
    data_noisy2 = data2 + 0.2*np.random.normal(size=data2.shape)
    
    initial_guess3 = (3, 3, 3, 100, 100, -100, -100, -100, 100, 20, 40, 20, 40, 20, 40)#, 0, 10)
    data_noisy3 = data3 + 0.2*np.random.normal(size=data2.shape)
    
    ### Choose fit 1 circle or 2 circle model
    plt.close('all')    
    
#    ### 1 circle model
#    fit_2D_Gaussian((x, y), data_noisy, initial_guess, twoD_Gaussian, plot=True)    
#    ### 2 circle model
#    fit_2D_Gaussian((x, y), data_noisy + data_noisy2, initial_guess2, two_blob, plot=True)


    fit_2D_Gaussian(xy, data_noisy, initial_guess, twoD_Gaussian, plot=True)    
    ### 2 circle model
    fit_2D_Gaussian(xy, data_noisy + data_noisy2, initial_guess2, two_blob, plot=True)
    ### 3 circle model
    fit_2D_Gaussian(xy, data_noisy + data_noisy2 + data_noisy3, initial_guess3, three_blob, plot=True)
















































